import os
import h5py  # 用于读写HDF5格式文件（高效存储大型数据集）
import torch
from torch.utils.data import Dataset, DataLoader  #  PyTorch数据集基类和数据加载器


class AVEDataset(Dataset):
    """
    AVE数据集类：继承PyTorch的Dataset基类，用于加载音频-视频多模态数据集
    支持读取视觉特征、音频特征及对应的标签，适配训练/验证/测试等不同数据拆分
    """
    def __init__(self, data_root, split='train'):
        """
        初始化数据集
        Args:
            data_root: 数据集根目录路径（包含所有HDF5格式的特征和标签文件）
            split: 数据拆分类型（'train'/'val'/'test'等，默认'train'），用于指定加载哪部分数据
        """
        super(AVEDataset, self).__init__()
        self.split = split  # 保存数据拆分类型
        # 定义各类数据文件的路径
        self.visual_feature_path = os.path.join(data_root, 'visual_feature.h5')  # 视觉特征HDF5文件路径
        self.audio_feature_path = os.path.join(data_root, 'audio_feature.h5')    # 音频特征HDF5文件路径
        self.labels_path = os.path.join(data_root, 'labels.h5')                  # 标签HDF5文件路径
        self.sample_order_path = os.path.join(data_root, f'{split}_order.h5')    # 当前拆分的数据索引文件路径（如train_order.h5）
        self.h5_isOpen = False  # HDF5文件是否已打开的标志（避免重复打开）

    def __getitem__(self, index):
        """
        按索引获取单一样本数据（PyTorch Dataset核心方法）
        Args:
            index: 样本在数据集中的索引（0-based）
        Returns:
            visual_feat: 视觉特征（numpy数组）
            audio_feat: 音频特征（numpy数组）
            label: 样本标签（numpy数组）
        """
        # 延迟打开HDF5文件：首次访问时才打开，节省内存
        if not self.h5_isOpen:
            # 打开视觉特征文件，读取'avadataset'数据集（HDF5中的数据集名称）
            self.visual_feature = h5py.File(self.visual_feature_path, 'r')['avadataset']
            # 打开音频特征文件，读取'avadataset'数据集
            self.audio_feature = h5py.File(self.audio_feature_path, 'r')['avadataset']
            # 打开标签文件，读取'avadataset'数据集
            self.labels = h5py.File(self.labels_path, 'r')['avadataset']
            # 打开当前拆分的索引文件，读取'order'数据集（存储当前拆分的样本索引）
            self.sample_order = h5py.File(self.sample_order_path, 'r')['order']
            self.h5_isOpen = True  # 标记文件已打开

        # 获取当前索引对应的样本在原始数据中的位置（通过split_order映射）
        sample_index = self.sample_order[index]
        # 根据原始索引提取视觉特征
        visual_feat = self.visual_feature[sample_index]
        # 根据原始索引提取音频特征
        audio_feat = self.audio_feature[sample_index]
        # 根据原始索引提取标签
        label = self.labels[sample_index]

        return visual_feat, audio_feat, label

    def __len__(self):
        """
        获取数据集样本总数（PyTorch Dataset核心方法）
        Returns:
            sample_num: 当前拆分的样本数量
        """
        # 打开当前拆分的索引文件，读取样本数量（通过'order'数据集的长度）
        f = h5py.File(self.sample_order_path, 'r')
        sample_num = len(f['order'])  # 'order'数据集的长度即样本数
        f.close()  # 读取后关闭文件，避免资源占用
        return sample_num